"""
Author: runyu.ma

flatten_loop_pass.py
"""

from ..core.dslc_parser import BinaryOp, BlockOp, ASTNode, Identifier, Number, String, UnaryOp
from ..core.dslc_symbol_table import SymbolTable
from ..core.dslc_misc import SafeDict

RECORDED_LOOP_CNT = 0

def flatten_loop_block(stab: SymbolTable, top: BlockOp):
    oplist = []
    local_val = SafeDict()
    global_val = {}
    #
    init_cond = top.oplist[1].oplist[0]
    if isinstance(init_cond, BinaryOp):
        init_expr = '{} = {}'.format(init_cond.l.name, init_cond.r.value)
    elif isinstance(init_cond, UnaryOp):
        init_expr = '{} = {}'.format(init_cond.x.oplist[0].l.name, init_cond.x.oplist[0].r.value)
    else:
        raise ValueError("for loop init cond illegal")
    #print("INIT FOR LOOP EXPR", init_cond, init_expr)
    exec(init_expr, globals=global_val, locals=local_val)
    #print(local_val)
    #
    exit_cond = top.oplist[1].oplist[1]
    #print(stab.lsvisiable())
    #print(exit_cond)
    if isinstance(exit_cond.r, Identifier):
        _param =  stab.lsvisiable()
        #support cnt{i} in for loop exit condition
        exit_cond.r.name = exit_cond.r.name.format_map(_param)
        if exit_cond.r.name in _param:
            exit_r_value = _param[exit_cond.r.name]
        else:
            print("flatten_loop corrupt, skip {} param {}".format(str(exit_cond), exit_cond.r.name))
            return top
    else:
        exit_r_value = exit_cond.r.value
    exit_expr = 'esig = ({} {} {})'.format(exit_cond.l.name, exit_cond.op, exit_r_value)
    #print("EXIT FOR LOOP EXPR", exit_cond, exit_expr)
    exec(exit_expr, globals=global_val, locals=local_val)
    #
    step_cond = top.oplist[1].oplist[2]
    step_expr = __flatten_binop(step_cond)
    #print("STEP FOR LOOP EXPR", step_cond, step_expr)
    #
    loop_body = top.oplist[2]
    #
    #在输出的BlockOp中增加一些Tag？？
    global RECORDED_LOOP_CNT
    RECORDED_LOOP_CNT += 1
    oplist.append(UnaryOp(':', Identifier('loop_no.{}'.format(RECORDED_LOOP_CNT))))
    #替换Identifier中name相同的
    #oplist.append(init_cond)
    loop_num = 0
    while local_val['esig']:
        loop_num += 1
        _content_block = []
        _content_block.append(UnaryOp(':', Identifier('repeat_no.{}'.format(loop_num))))
        _content_block.extend(BinaryOp('=', Identifier('{}_{}'.format(e, local_val[e])), Number(local_val[e])) for e in local_val if e != 'esig')
        _content_block.extend(__replace_loopidx(loop_body, local_val).oplist)
        print(_content_block)
        #
        oplist.append(BlockOp(_content_block))
        print(step_expr)
        exec(step_expr, globals=global_val, locals=local_val)
        exec(exit_expr, globals=global_val, locals=local_val)
        print(local_val)
    return BlockOp(oplist)


def __flatten_binop(bop: ASTNode):
    if isinstance(bop, Identifier):
        return bop.name
    if isinstance(bop, Number):
        return bop.value
    if isinstance(bop, BinaryOp):
        return '{} {} {}'.format(__flatten_binop(bop.l), bop.op, __flatten_binop(bop.r))
    raise ValueError("__flatten_binop can not parse Op")


def __replace_loopidx(top: ASTNode, local_val: dict):
    #oplist = []
    #for e in top.oplist:
    if isinstance(top, BlockOp):
        oplist = []
        for e in top.oplist:
            oplist.append(__replace_loopidx(e, local_val))
        return BlockOp(oplist)
    elif isinstance(top, Identifier) and top.name in local_val:
        return Number(local_val[top.name])
    elif isinstance(top, String):# and isinstance(top.parent, BlockOp):
        newstr = str(top.value)
        for e in local_val:
            #newstr = newstr.replace('{{{}}}'.format(e), '{{{}_{}}}'.format(e, local_val[e]))
            newstr = newstr.replace('{{{}}}'.format(e), str(local_val[e]))
        return String(newstr)
    elif isinstance(top, Identifier):
        newstr = str(top.name)
        for e in local_val:
            newstr = newstr.replace('{{{}}}'.format(e), '{}'.format(local_val[e]))
            #newstr = newstr.replace('{{{}}}'.format(e), '{{{}_{}}}'.format(e, local_val[e]))
        return Identifier(newstr)
    elif isinstance(top, UnaryOp) and isinstance(top.x, BinaryOp):
        e2 = top.copy()
        if isinstance(e2.x.r, String):
            e2.x.r.value = e2.x.r.value.format_map(local_val)
        return e2
    elif isinstance(top, UnaryOp):
        return UnaryOp(top.op, __replace_loopidx(top.x, local_val))
    elif isinstance(top, BinaryOp):
        return BinaryOp(top.op, __replace_loopidx(top.l, local_val), __replace_loopidx(top.r, local_val))
    else:
        return top.copy()
    
